
#include "../include/SSS_SOLVE.h"
//#include <__clang_cuda_device_functions.h>
#include <cuda_runtime_api.h>
#include <cusparse.h>
#include <cuda_runtime.h>
#include <driver_types.h>
#include "../../Main/include/CUDA_error.cuh"

#include <helper_cuda.h>

#include <nsparse.cuh>
#include <nsparse_asm.h>

//#include <>


double pre_smooth_time = 0;
double post_smooth_time = 0;
double residual_time = 0;
double restriction_time = 0;
double prolongation_time = 0;
double direct_LU_time = 0;
double norm2_time = 0;
double cuda_malloc_memcpy_time = 0;
double cusparse_init_time = 0;
double cublas_time = 0;
double csr2dense_time = 0 ;

extern double SpMV_count ;

void SSS_Solve_print(double solve_time,double pre_smooth_time,double post_smooth_time,double residual_time,double restriction_time,double prolongation_time,double direct_LU_time,double norm2_time,double cuda_malloc_memcpy_time,double cusparse_init_time,double cublas_time , double csr2dense_time )
{
    printf("\n");
    printf("solve_time            = %.2lf ms          \n",solve_time);
    printf("---------------------Real Time: ---------------------\n");
    printf("pre_smooth_time       = %.2lf ms          \n",pre_smooth_time);
    printf("post_smooth_time      = %.2lf ms          \n",post_smooth_time);
    printf("residual_time         = %.2lf ms          \n",residual_time);
    printf("restriction_time      = %.2lf ms          \n",restriction_time);
    printf("prolongation_time     = %.2lf ms          \n",prolongation_time);
    printf("direct_LU_time        = %.2lf ms          \n",direct_LU_time);
    printf("norm2_time            = %.2lf ms          \n",norm2_time);
    printf("cuda_malloc_cpy_time  = %.2lf ms          \n",cuda_malloc_memcpy_time);
    printf("cusparse_init_time    = %.2lf ms          \n",cusparse_init_time);
    printf("cublas_time           = %.2lf ms          \n",cublas_time);
    printf("csr2dense_time        = %.2lf ms          \n",csr2dense_time);

    double pre = pre_smooth_time/solve_time *100;
    double post = post_smooth_time/solve_time * 100;
    double r = residual_time/solve_time * 100;
    double res = restriction_time/solve_time * 100;
    double pro = prolongation_time/solve_time * 100;
    double direct = direct_LU_time/solve_time *100;
    double norm2 = norm2_time/solve_time * 100;
    double cuda_malloc_cpy_ratio = cuda_malloc_memcpy_time/solve_time * 100;
    double cusparse = cusparse_init_time/solve_time * 100;
    double cublas = cublas_time/solve_time * 100;
    double c2d = csr2dense_time/solve_time * 100;

    
    printf("\n");
    printf("---------------------Component/Solve ratio: ---------------------\n");
    printf("pre_smooth      /  Solve   = %.2lf %%      \n",pre);
    printf("post_smooth     /  Solve   = %.2lf %%      \n",post);
    printf("residual        /  Solve   = %.2lf %%      \n",r);
    printf("restriction     /  Solve   = %.2lf %%      \n",res);
    printf("prolongation    /  Solve   = %.2lf %%      \n",pro);
    printf("direct_LU       /  Solve   = %.2lf %%      \n",direct);
    printf("norm2           /  Solve   = %.2lf %%      \n",norm2);
    printf("cuda_malloc_cpy /  Solve   = %.2lf %%      \n",cuda_malloc_cpy_ratio);
    printf("cusparse_init   /  Solve   = %.2lf %%      \n",cusparse);
    printf("cublas_init     /  Solve   = %.2lf %%      \n",cublas);
    printf("c2d             /  Solve   = %.2lf %%      \n",c2d);



    printf("\n");
}





void scr2sfCSR2(sfCSR *csr_mat,SSS_MAT *mat)
{
        csr_mat->M = mat->num_rows;
        csr_mat->N = mat->num_cols;
        csr_mat->nnz = mat->num_nnzs;

        checkCudaErrors(cudaMalloc((void **)&(csr_mat->d_rpt), sizeof(int) * (csr_mat->M + 1)));
        checkCudaErrors(cudaMalloc((void **)&(csr_mat->d_col), sizeof(int) * csr_mat->nnz));
        checkCudaErrors(cudaMalloc((void **)&(csr_mat->d_val), sizeof(real) * csr_mat->nnz));
    
        checkCudaErrors(cudaMemcpy(csr_mat->d_rpt, mat->d_row_ptr, sizeof(int) * (csr_mat->M + 1), cudaMemcpyDeviceToDevice));
        checkCudaErrors(cudaMemcpy(csr_mat->d_col, mat->d_col_idx, sizeof(int) * csr_mat->nnz, cudaMemcpyDeviceToDevice));
        checkCudaErrors(cudaMemcpy(csr_mat->d_val, mat->d_val, sizeof(real) * csr_mat->nnz, cudaMemcpyDeviceToDevice));
}

void spmv_amb2(sfCSR *csr_mat, real *x, real *y, sfPlan *plan)
{
    real *d_x, *d_y;
    sfAMB mat;


    checkCudaErrors(cudaMalloc((void **)&d_x, sizeof(real) * (csr_mat->N + MAX_BLOCK_SIZE)));
    checkCudaErrors(cudaMalloc((void **)&d_y, sizeof(real) * (csr_mat->M + WARP)));
    checkCudaErrors(cudaMemcpy(d_x, x, sizeof(real) * csr_mat->N, cudaMemcpyDeviceToDevice));
  
    cudaDeviceSynchronize();
    sf_csr2amb(&mat, csr_mat, d_x, plan);
    cudaDeviceSynchronize();

  
    /* Execution of SpMV on Device */
    sf_spmv_amb(d_y, &mat, d_x, plan);
    cudaDeviceSynchronize();

    checkCudaErrors(cudaMemcpy(y, d_y, sizeof(real) * csr_mat->M, cudaMemcpyDeviceToDevice));
  
    cudaDeviceSynchronize();

    /*7th step. Release*/
    cudaFree(d_x);
    cudaFree(d_y);
    release_amb(mat);
    //release_csr(*csr_mat);
  
}




SSS_RTN SSS_amg_solve_cuda(SSS_AMG *mg, SSS_VEC *x, SSS_VEC *b)
{

    SSS_MAT *ptrA;
    SSS_VEC *r;

    double alpha  = 1.0;
    double n_alpha = -1.0;
    double beta  = 0.0;
    double beta_1  = 1.0f;

    int max_it;
    double tol; //  容差，达到tol就停止
    double sumb; // L2norm(b) 

    double solve_start, solve_end; //计时
    double relres1 = 1.0, absres0, absres, factor;  //相对残差   初始残差  残差   收缩因子
    int iter = 0;
    SSS_RTN rtn;
    
    solve_start = SSS_get_time(); 

    /* check pars */
    assert(mg != NULL);
    assert(x != NULL);
    assert(b != NULL);

    
    cudaError_t E;  

    //cublas handle
    double cublas_time1  = SSS_get_time();
    cublasHandle_t cublas_hndl_solve;
    cublasStatus_t cublas_stat_solve;
    cublas_stat_solve = cublasCreate(&cublas_hndl_solve);
    double cublas_time2  = SSS_get_time();
    double cublas_time  = cublas_time2 - cublas_time1;
    
    // cusparse handle
    double cusparse_init1 = SSS_get_time();

    cusparseHandle_t cusparseHandle = NULL;   
    cusparseStatus_t cusparse_stat;
 
    cusparseSpMatDescr_t matA;
    cusparseDnVecDescr_t vecX, vecY;
    void*                dBuffer    = NULL;
    size_t               bufferSize = 0;


    cusparse_stat  = cusparseCreate(&cusparseHandle);

    // cusparseMatDescr_t descr = 0;
    // cusparseCreateMatDescr(&descr);
    double cusparse_init2 = SSS_get_time();

    cusparse_init_time+=cusparse_init2 - cusparse_init1;

    max_it = mg->pars.max_it;
    tol = mg->pars.tol;

    ptrA = &mg->cg[0].A;
    r = &mg->cg[0].wp;

   

    //计算b的二范数
    double norm2_1 = SSS_get_time();
    cublas_stat_solve = cublasDnrm2(cublas_hndl_solve,b->n,b->d_d,1,&sumb);
    double norm2_2 = SSS_get_time();
    norm2_time += norm2_2 - norm2_1;

    absres0 = sumb;     //初始残差等于b的二范数


    // Print iteration information if needed
    SSS_print_itinfo( STOP_REL_RES, iter, 1.0, sumb, 0.0);

    /* init, to make compiler hAppy */
    rtn.ares = 0;
    rtn.rres = 0;
    rtn.nits = 0;

    if (fabs(sumb) == 0.) 
    {
        SSS_vec_set_value(x, 0);
        mg->rtn = rtn;
        return rtn;
    }

    //将初始x的数据拷贝到gpu
    double cuda_api3 = SSS_get_time();
    E = cudaMalloc((void **)&x->d_d,(x->n) * sizeof(double));
    E = cudaMemcpy(x->d_d, x->d, (x->n) * sizeof(double), cudaMemcpyHostToDevice);
    double cuda_api4 = SSS_get_time();
    cuda_malloc_memcpy_time += cuda_api4 - cuda_api3;


    /* set x and b */
    mg->cg[0].x = *x;
    mg->cg[0].b = *b;
   
    /* set x and b on gpu */
    //直接将gpu上x的数据拷贝给mg->cg[0].x
    
    //x
    double cuda_api5 = SSS_get_time();
    E = cudaMemcpy(mg->cg[0].x.d_d, x->d_d, (x->n) * sizeof(double), cudaMemcpyDeviceToDevice);
    //b
    E = cudaMemcpy(mg->cg[0].b.d_d, b->d_d, (b->n) * sizeof(double), cudaMemcpyDeviceToDevice);
    double cuda_api6 = SSS_get_time();
    cuda_malloc_memcpy_time += cuda_api6 - cuda_api5;


    //select spmv 
    int method = 0;   //method = 0,cusparse,      method = 1,csr5   method = 2,nsparse

    //inti nsparse
    sfCSR mat_A;
    scr2sfCSR2(&mat_A,&mg->cg[0].A);
    sfPlan plan;
    init_plan(&plan);
    
    while ((++iter <= max_it)) 
    {    
        SSS_amg_cycle_cuda(mg);
        double residual_1;
        if (method == 0)
        {
            printf("cusparse spmv \n");
            double cuda_api13 = SSS_get_time();
            E = cudaMemcpy(r->d_d, b->d_d, (r->n) * sizeof(double), cudaMemcpyDeviceToDevice);
            cudaDeviceSynchronize();
            
            double cuda_api14 = SSS_get_time();
            cuda_malloc_memcpy_time += cuda_api14 - cuda_api13;


            // Form residual r = b - A*x  
            residual_1 = SSS_get_time();

            //用cusparse的SpMV算 r = alpha(n_alpha,-1) * Ax + beta(beta_1,1) * r(b)
            // r = -AX + r
            // Create sparse matrix A in CSR format
            cusparseCreateCsr(&matA, ptrA->num_rows, ptrA->num_cols,ptrA->num_nnzs,
                                        ptrA->d_row_ptr,ptrA->d_col_idx, ptrA->d_val,
                                        CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I,
                                        CUSPARSE_INDEX_BASE_ZERO, CUDA_R_64F) ;
            // Create dense vector X
            cusparseCreateDnVec(&vecX, x->n, x->d_d, CUDA_R_64F);
            // Create dense vector y
            cusparseCreateDnVec(&vecY, r->n, r->d_d, CUDA_R_64F);
            // allocate an external buffer if needed
            cusparseSpMV_bufferSize(cusparseHandle, CUSPARSE_OPERATION_NON_TRANSPOSE,
                                    &n_alpha, matA, vecX, &beta_1, vecY, CUDA_R_64F,
                                    CUSPARSE_MV_ALG_DEFAULT, &bufferSize);
            
            cudaMalloc(&dBuffer, bufferSize);
            // execute SpMV
            cusparse_stat = cusparseSpMV(cusparseHandle, CUSPARSE_OPERATION_NON_TRANSPOSE,
                        &n_alpha, matA, vecX, &beta_1, vecY, CUDA_R_64F,
                        CUSPARSE_MV_ALG_DEFAULT, dBuffer);
        }

        //nsparse spmv

        else if(method==2)
        {
            printf("nsparse spmv \n");
            residual_1 = SSS_get_time();
            E = cudaMemcpy(r->d_d, b->d_d, (r->n) * sizeof(double), cudaMemcpyDeviceToDevice);
            cudaDeviceSynchronize();
            
            //nsparse spmv
            //(r->d_d = Ax)
            double *dtmp;
            E = cudaMalloc((void **)&dtmp,(r->n) * sizeof(double));
            cudaDeviceSynchronize();

            spmv_amb2(&mat_A, x->d_d, dtmp, &plan);
            cudaDeviceSynchronize();
            // y(b->d_d) = alpha(n_alpha,-1)*x(r->d_d,Ax) +y(b->d_d)
            //b = -Ax +b
            cublas_stat_solve = cublasDaxpy_v2(cublas_hndl_solve, ptrA->num_rows, &n_alpha, dtmp, 1, r->d_d, 1); //实现向量+
            cudaDeviceSynchronize();
            //再用cublas的axpy算 b - Ax 
            if(cublas_stat_solve !=CUBLAS_STATUS_SUCCESS)
            {
                printf("cublas_stat_solve !!!\n");
            }
        }
        SpMV_count++;


        cudaDeviceSynchronize();
        //cublas_stat_solve = cublasDaxpy_v2(cublas_hndl_solve, ptrA->num_rows, &n_alpha, b->d_d, 1, r->d_d, 1); //实现向量+
        // cublas_stat_solve = cublasDaxpy_v2(cublas_hndl_solve, ptrA->num_rows, &alpha, r->d_d, 1, b->d_d, 1); //实现向量+

        //alpha_spmv_cuda(alpha, ptrA->num_rows, ptrA->d_row_ptr, ptrA->d_col_idx, ptrA->d_val, x->d_d, r->d_d);

        cudaDeviceSynchronize();
        double residual_2 = SSS_get_time();
        residual_time += residual_2 - residual_1;

        // Compute norms of r and convergence factor 
        //absres 取r（residual ）的norm2 ||r||
        double norm2_3 = SSS_get_time();
        cublas_stat_solve = cublasDnrm2(cublas_hndl_solve,r->n,r->d_d,1,&absres);
        // cublas_stat_solve = cublasDnrm2(cublas_hndl_solve,b->n,b->d_d,1,&absres);
        double norm2_4 = SSS_get_time();
        norm2_time += norm2_4 - norm2_3;

        /*相对残差*/
        relres1 = absres / sumb;        // relative residual ||r||/||b||
        //收敛因子
        factor = absres / absres0;      // contraction factor  
        
        absres0 = absres;               // prepare for next iteration

        // Print iteration information if needed
        SSS_print_itinfo( STOP_REL_RES, iter, relres1, absres, factor);
        
        /* save convergence info */
        rtn.ares = absres;
        rtn.rres = relres1;
        rtn.nits = iter;
        mg->rtn = rtn;
        // Check convergence
        if (relres1 < tol) break;   //两种退出条件， iter >= max_iter    或  relres1 < tol
    }
    release_csr(mat_A);

    // cusparseDestroySpMat(matA);
    // cusparseDestroyDnVec(vecX);
    // cusparseDestroyDnVec(vecY);
    // cusparseDestroy(cusparseHandle);


    solve_end = SSS_get_time();
    solve_time = solve_end - solve_start;
    //printf solve time
    SSS_Solve_print(solve_time,pre_smooth_time,post_smooth_time,residual_time,restriction_time,prolongation_time,direct_LU_time,norm2_time,cuda_malloc_memcpy_time,cusparse_init_time,cublas_time,csr2dense_time);



    return rtn;
}